import argparse
import os
from vortex.engine.train import Trainer
from vortex.engine.inference import ModelInference
from vortex.utils.config import get_config, get_config_subs


def train(args):
    config_path = args.config
    # configs = get_config_subs(config_path)
    configs = get_config(config_path)

    trainer = Trainer(configs)
    trainer.train()


def inference(args):
    config_path = args.config
    configs = get_config(config_path)

    inference = ModelInference(configs)

    image_path = configs.INFER.IMAGE_PATH

    if os.path.isfile(image_path):
        inference.detect(image_path)
    elif os.path.isdir(image_path):
        for image_name in os.listdir(image_path):
            if image_name.endswith('.jpg'):
                image_path = os.path.join(image_path, image_name)
                inference.detect(image_path)
    else:
        raise ValueError("invalid image path")


def parse_args():
    parser = argparse.ArgumentParser(description='Simple Object Detection')
    parser.add_argument('--task', default='train', help='train/test/inference')
    parser.add_argument('--config', help='configuration path')

    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    if args.task == 'train':
        train(args)
    elif args.task == 'inference':
        inference(args)


if __name__ == '__main__':
    main()
